%% osl_ica_preproc.m
%
%  S_out=osl_ica_preproc(S_in)
%  
%  Input structure S_in contains:
%     REQUIRED: source_recon_results_fname - the name of the source
%               recontructed data output by the beamformer.
%               oil - the oil structure that defines the analysis
%               parameters.
%     OPTIONAL: window_length            : the window over which the envelope is downsampled (default 1s).   
%               overlap                  : the overlap between windows (default is 0).
%               timewindow               : subset of trial window to remove edge effects (default is 'all').
%               ss                       : spatial smoothing of down-sampled envelopes by gaussian kernel (FWHM) (default is 4mm).               
%               ds                       : spatial down-sampling of data (default 8mm).
%
%  Output structure S_out contains:      
%               fils_nifti           : name of the nifti file output 
%               fils_nifti_nonorm    : name of the nifti file output for
%                                      the un-normalised data.

%
%  HL 060213
%  Version 1.2

function S_out=osl_ica_preproc(S_in)

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Initialise Variables

global OSLDIR;

if isfield(S_in,'source_recon_results_fname'),    source_recon_results_fname  =  S_in.source_recon_results_fname;  else     error('source_recon file name not specified');                                                    end;
if isfield(S_in,'oil'),                           oil                         =  S_in.oil;                         else     error('oil not provided');                                                                        end;
if isfield(S_in,'window_length'),                 window_length               =  S_in.window_length;               else     warning('window_length not specified, Set to 1s by default');   window_length            = 1;     end;
if isfield(S_in,'ss'),                            ss                          =  S_in.ss;                          else     warning('ss not specified, Set to 4mm by default');             ss                       = 4;     end;
if isfield(S_in,'ds'),                            ds                          =  S_in.ds;                          else     warning('ds not specified, Set to 8mm by default');             ds                       = 8;     end;
if isfield(S_in,'overlap'),                       overlap                     =  S_in.overlap;                     else     warning('overlap not specified, Set to 0 by default');          overlap                  = 0;     end;
if isfield(S_in,'timewindow'),                    timewindow                  =  S_in.timewindow;                  else     warning('timewindow not specified, Using whole trial window,'); timewindow               = 'all'; end;

save_dir = [oil.source_recon.dirname '/' oil.enveloping.name];

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Load Beamformed Data

source_recon_results=osl_load_oil_results(oil,source_recon_results_fname);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Construction of Source Space Data

D=source_recon_results.BF.data.D;
Ntrials=0;
for i=1:length(source_recon_results.source_recon.conditions), % indexes conditions
        trigname=source_recon_results.source_recon.conditions{i};
        Ntrials=Ntrials+length(D.pickconditions(trigname)); %% number of trials for this condition
end

if Ntrials==1 ;
    data_type='continuous';
else    
    data_type='epoched';
end

Nsamples=D.nsamples;
Nvoxels=size(source_recon_results.BF.sources.pos,1);

modality='MEG';
chanindmeg = strmatch(modality, D.chantype);
chanind = setdiff(chanindmeg, D.badchannels);

ft_progress('init', 'etf');

switch data_type
    case 'continuous'
        fprintf('\n%s%s%s\n','Processing ',data_type, ' data. Converting to time domain, enveloping and down-sampling')
        % Continuous Data
        % Select data of interest
        if isempty(oil.source_recon.time_range),
            time_range = [D.time(1) D.time(end)];
        else
            time_range = oil.source_recon.time_range;
        end
        Events = D.events; GoodEpochs=[D.time(1) D.time(end)];
        if D.ntrials==1
            for ev = 1:numel(Events)
                if strcmp(Events(ev).type,'BadEpoch')
                    GoodEpochs(end,2)=Events(ev).time;
                    GoodEpochs(end+1,1)=Events(ev).time+Events(ev).duration;
                    GoodEpochs(end,2)=D.time(end);
                end
            end
        end

        good_samples=zeros(1,D.nsamples);
        for i=1:size(GoodEpochs,1)
            good_samples(D.indsample(GoodEpochs(i, 1)):D.indsample(GoodEpochs(i, 2)))=1;
        end
        samples_of_interest=zeros(1,D.nsamples);
        for i=1:size(time_range,1)
            samples_of_interest(D.indsample(time_range(i, 1)):D.indsample(time_range(i, 2)))=1;
        end

        samples2use = samples_of_interest & good_samples;

        % Set size of ica data based on length of data excluding bad
        % sections
        winsize  = window_length*D.fsample;
        t        = D.time;
        t(~samples2use) = [];
        ica_course_nonorm = zeros(Nvoxels, length(osl_movavg(t,t,winsize,overlap,0)));
        ica_course = zeros(size(ica_course_nonorm));
        var_dat = zeros(Nvoxels,1);
        sensor_data=D(chanind, :,1);
        
        % Loop over voxels computing averaged envelope
        for vox=1:Nvoxels 
            ft_progress(vox/Nvoxels);
            weights=source_recon_results.BF.inverse.W.MEG{1}{vox};
            dat=weights*sensor_data;
            var_dat(vox) = var(dat(samples2use))/(weights*weights');
            % Hilbert Envelope
            HE = abs(hilbert(dat));
            HE(~samples2use) = [];
            ica_course_nonorm(vox,:) = osl_movavg(HE,t,winsize,overlap,0);
            ica_course(vox,:) = ica_course_nonorm(vox,:)/sqrt(weights*weights');
        end
        
    case 'epoched'
        % Epoched Data 
        % Epoches are kept in the order they were recorded originally.
        fprintf('\n%s%s%s\n','Processing ',data_type, ' data. Converting to time domain, enveloping and down-sampling')
        
        winsize  = window_length*D.fsample;
        t        = D.time;
        
        % select subset of times to use
        if strcmp(timewindow,'all'), % use times specified at source recon level
            isSourceReconTimeSpecified = ...
                (isfield(oil.source_recon, 'time_range') && ...
                 isequal(length(oil.source_recon.time_range), 2));
             
            if isSourceReconTimeSpecified,
                timewindow = [max(t(1), oil.source_recon.time_range(1)), ...
                              min(t(end), oil.source_recon.time_range(2))]; 
            else
                timewindow = [t(1), t(end)];
            end%if
        else
            % timewindow should be a 2-component vector of times to use
            if ~isnumeric(timewindow) || ~isequal(length(timewindow), 2),
                error([mfilename ':WrongTimewindowFormat'], ...
                      ['timewindow must be ''all'' ', ...
                       'or a 2-component vector. \n']);
            end%if                      
        end%if
        time_to_use = (t>=timewindow(1) & t<=timewindow(2));
        
        % declare memory
        ica_course_nonorm = zeros(Nvoxels, Ntrials, length(osl_movavg(t(time_to_use),t(time_to_use),winsize,overlap,0)));
        ica_course        = zeros(size(ica_course_nonorm));
        sensor_data       = D(chanind, :,:);
        var_dat           = zeros(Nvoxels,Ntrials);

        % perform enveloping and downsample
        for vox = 1:Nvoxels, % indexes brain space
            weights = source_recon_results.BF.inverse.W.MEG{1}{vox};
            
            for tri = 1:Ntrials
                dat = weights*sensor_data(:, :, tri);
                var_dat(vox,tri) = var(dat(time_to_use))/(weights*weights');
                
                HE = abs(hilbert(dat));
                
                ica_course_nonorm(vox,tri,:) = osl_movavg(HE(time_to_use), ...
                                                          t(time_to_use), ...
                                                          winsize, ...
                                                          overlap, ...
                                                          0);
                %ica_course(vox,tri,:)        = ica_course_nonorm/sqrt(weights*weights');
                ica_course(vox,tri,:)        = ica_course_nonorm(vox,tri,:) ...
                                                / sqrt(weights * weights'); % MW fix from astle/henry emails
            end;
        end       
        var_dat = mean(var_dat,2);

        % Remove trial structure
        c=1; Ntpts=size(ica_course,3);
        ica_course_rs = zeros([size(ica_course,1), size(ica_course,2)*size(ica_course,3)]);
        ica_course_nonorm_rs = zeros([size(ica_course,1), size(ica_course,2)*size(ica_course,3)]);
        for i=1:size(ica_course,2)
            ica_course_rs(:,c:c+Ntpts-1)=permute(ica_course(:,i,:),[1 3 2]);
            ica_course_nonorm_rs(:,c:c+Ntpts-1)=permute(ica_course_nonorm(:,i,:),[1 3 2]);
            c=c+Ntpts;
        end
        ica_course        = ica_course_rs;
        ica_course_nonorm = ica_course_nonorm_rs;
        clear ica_course_rs ica_course_nonorm_rs;
    otherwise
end

ft_progress('close');

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Convert to MNI nii space and save nii
try    
    fprintf('\n%s%1.1f%s%1.1f%s\n','Saving Downsampled Envelopes as .nii files and applying ', ss, 'mm spatial smoothing and ', ds, 'mm spatial resampling.')
    fnamec=[save_dir '/' source_recon_results_fname '_winavHE_delta_' num2str(window_length) 's'];
    nii_quicksave(ica_course,fnamec,source_recon_results.gridstep);
    
    fnamecnn=[save_dir '/' source_recon_results_fname '_NoWeightsNorm_winavHE_delta_' num2str(window_length) 's'];
    nii_quicksave(ica_course_nonorm,fnamecnn,source_recon_results.gridstep);
    
    if ~isdir([save_dir '/variance_maps/']);mkdir([save_dir '/variance_maps/']); end
    nii_quicksave(var_dat,[save_dir '/variance_maps/' source_recon_results_fname '_noise_corrected_variance_map'],source_recon_results.gridstep,2);

    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
    %% spatial smoothing and downsample,
    
    % Weights Normalised Data
    fnamec2=[fnamec '_ss' num2str(ss) 'mm'];
    maskimage=[OSLDIR '/std_masks/MNI152_T1_' num2str(source_recon_results.gridstep) 'mm_brain_mask'];
    
    runcmd(['fslmaths ' fnamec  ' -s ' num2str(ss) ' -mas ' maskimage ' tmp1']);
    runcmd(['fslmaths ' maskimage ' -s ' num2str(ss)  ' -mas ' maskimage ' tmp2']);
    runcmd(['fslmaths tmp1  -div tmp2 '  fnamec2]);
    runcmd('rm tmp1.nii.gz tmp2.nii.gz');
    
    fnamec3=[fnamec2 '_ds' num2str(ds) 'mm'];
    runcmd(['flirt -in ' fnamec2 ' -applyxfm -init ' getenv('FSLDIR') '/etc/flirtsch/ident.mat -out ' fnamec3 ...
        ' -paddingsize 0.0 -interp trilinear -ref ' OSLDIR '/std_masks/MNI152_T1_' num2str(ds) 'mm_brain.nii.gz']);         % Spatial Downsampling
    runcmd(['fslmaths ' fnamec3 ' -mas ' OSLDIR '/std_masks/MNI152_T1_' num2str(ds) 'mm_brain_mask ' fnamec3]);             % mask to reduce edge blurring
    
    [~, nam, ext] = fileparts(fnamec3);
    S_out.fils_nifti=[nam ext];
    
    % Weights Un-Normalised Data
    fnamecnn2=[fnamecnn '_ss' num2str(ss) 'mm'];
    maskimage=[OSLDIR '/std_masks/MNI152_T1_' num2str(source_recon_results.gridstep) 'mm_brain_mask'];
    
    runcmd(['fslmaths ' fnamecnn  ' -s ' num2str(ss) ' -mas ' maskimage ' tmp1']);
    runcmd(['fslmaths ' maskimage ' -s ' num2str(ss)  ' -mas ' maskimage ' tmp2']);
    runcmd(['fslmaths tmp1  -div tmp2 '  fnamecnn2]);
    runcmd('rm tmp1.nii.gz tmp2.nii.gz');
    
    fnamecnn3=[fnamecnn2 '_ds' num2str(ds) 'mm'];
    runcmd(['flirt -in ' fnamecnn2 ' -applyxfm -init ' getenv('FSLDIR') '/etc/flirtsch/ident.mat -out ' fnamecnn3 ...
        ' -paddingsize 0.0 -interp trilinear -ref ' OSLDIR '/std_masks/MNI152_T1_' num2str(ds) 'mm_brain.nii.gz']);         % Spatial Downsampling
    runcmd(['fslmaths ' fnamecnn3 ' -mas ' OSLDIR '/std_masks/MNI152_T1_' num2str(ds) 'mm_brain_mask ' fnamecnn3]);             % mask to reduce edge blurring
    
    [~, nam, ext] = fileparts(fnamecnn3);
    S_out.fils_nifti_nonorm=[nam ext];
        
catch
    fprintf('\n%s\n','Unable to save as .nii. Saving as .mat instead. No spatial smoothing or spatial resampling will be applied.')
    fnamec=[save_dir '/' source_recon_results_fname '_winavHE_delta_' num2str(window_length) 's' '_ss' num2str(ss) 'mm' '_ds' num2str(ds) 'mm'];
    save(fnamec,'ica_course');
    S_out.fils_nifti=fnamec;
    
    fnamecnn=[save_dir '/' source_recon_results_fname '_NoWeightsNorm__winavHE_delta_' num2str(window_length) 's' '_ss' num2str(ss) 'mm' '_ds' num2str(ds) 'mm'];
    save(fnamecnn,'ica_course_nonorm');
    S_out.fils_nifti_nonorm=fnamec;
end
